import subprocess
import argparse
import os

# Define a function to run the sample-bt tool with specified parameters
def run_sample_bt(output_file, pid=None, duration=None, user_space=False, extra_args=None):
    cmd = ['./sample-bt']
    if pid:
        cmd.extend(['-p', str(pid)])
    if duration:
        cmd.extend(['-t', str(duration)])
    if user_space:
        cmd.append('-u')
    if extra_args:
        cmd.extend(extra_args.split())
    cmd.extend(['>', output_file])

    print(f"Running command: {' '.join(cmd)}")
    subprocess.run(' '.join(cmd), shell=True, check=True)

# Define a function to generate a flame graph from the input file
def generate_flamegraph(input_file, output_file):
    folded_file = os.path.splitext(output_file)[0] + '.folded'
    
    cmd_stackcollapse = ['stackcollapse-stap.pl', input_file]
    with open(folded_file, 'w') as f:
        print(f"Running command: {' '.join(cmd_stackcollapse)}")
        subprocess.run(cmd_stackcollapse, stdout=f, check=True)

    cmd_flamegraph = ['flamegraph.pl', '-o', output_file]
    print(f"Running command: {' '.join(cmd_flamegraph)}")
    subprocess.run(cmd_flamegraph, stdin=open(folded_file), check=True)

# Main function that defines command line arguments and calls other functions to perform the task
def main():
    parser = argparse.ArgumentParser(description='Generate flame graph from SystemTap sampling.')
    parser.add_argument('-p', '--pid', type=int, help='Specify the user process PID.')
    parser.add_argument('-t', '--duration', type=int, help='Specify the number of seconds for sampling.')
    parser.add_argument('-u', '--user', action='store_true', help='Sample in the user-space.')
    parser.add_argument('-a', '--extra', type=str, help='Pass extra arguments to the sample-bt script.')
    parser.add_argument('-o', '--output', type=str, default='flamegraph.svg', help='Output flame graph file name.')

    args = parser.parse_args()

    output_bt = os.path.splitext(args.output)[0] + '.unfold'
    run_sample_bt(output_bt, args.pid, args.duration, args.user, args.extra)
    generate_flamegraph(output_bt, args.output)

if __name__ == '__main__':
    main()
